# -*- coding: utf-8 -*-

from fastapi import APIRouter, HTTPException, UploadFile, status
from models.OCRModel import *
from models.RestfulModel import *
from paddleocr import PaddleOCR
from utils.ImageHelper import base64_to_ndarray, bytes_to_ndarray
import requests
import os

router = APIRouter(prefix="/ocr", tags=["OCR"])

#旧的初始化方式
#OCR_LANGUAGE = os.environ.get("OCR_LANGUAGE", "ch")
#ocr = PaddleOCR(use_angle_cls=True, lang=OCR_LANGUAGE)

# 设置OCR参数
ocr_params = {
    "use_angle_cls": True,  # 是否使用角度分类器，用于文本检测
    "use_gpu": False,       # 是否使用GPU，需要在机器上有GPU支持
    "lang": 'ch',           # 识别的语言，中文设置为'ch'
    "drop_score": 0.1       # 设置这个参数无效
}
ocr = PaddleOCR(**ocr_params)


@router.get('/predict-by-path', response_model=RestfulModel, summary="识别本地图片")
def predict_by_path(image_path: str):
    result = ocr.ocr(image_path, cls=True)
    restfulModel = RestfulModel(success=True, message="识别本地图片成功", data=result_handle(result), cls=OCRModel)
    return restfulModel


@router.post('/predict-by-base64', response_model=RestfulModel, summary="识别 Base64 数据")
def predict_by_base64(base64model: Base64PostModel):
    img = base64_to_ndarray(base64model.base64_str)
    result = ocr.ocr(img=img, cls=True)
    restfulModel = RestfulModel(success=True, message="识别本地图片成功", data=result_handle(result), cls=OCRModel)
    return restfulModel


@router.post('/predict-by-file', response_model=RestfulModel, summary="识别上传文件")
async def predict_by_file(file: UploadFile):
    restfulModel: RestfulModel = RestfulModel()
    if file.filename.endswith((".jpg", ".png")):  # 只处理常见格式图片        
        file_data = file.file
        file_bytes = file_data.read()
        img = bytes_to_ndarray(file_bytes)
        result = ocr.ocr(img=img, cls=True)
        restfulModel.success = True
        restfulModel.message = "识别上传文件成功:[" + file.filename + "]"
        restfulModel.data = result_handle(result)
    else:
        restfulModel.success = False
        restfulModel.message = "识别上传文件失败,只支持[jpg]、[png]上传文件[" + file.filename + "]"        
    return restfulModel


@router.get('/predict-by-url', response_model=RestfulModel, summary="识别图片 URL")
async def predict_by_url(imageUrl: str):
    restfulModel: RestfulModel = RestfulModel()
    response = requests.get(imageUrl)
    image_bytes = response.content
    if image_bytes.startswith(b"\xff\xd8\xff") or image_bytes.startswith(b"\x89PNG\r\n\x1a\n"):  # 只处理常见格式图片 (jpg / png)         
        img = bytes_to_ndarray(image_bytes)
        result = ocr.ocr(img=img, cls=True)
        restfulModel.success = True
        restfulModel.data = result_handle(result)
        restfulModel.message = "识别URL成功"
    else:
        restfulModel.success = False
        restfulModel.message = "识别URL失败,只支持[jpg]、[png]上传文件[" + file.filename + "]"
    return restfulModel

def result_handle(list : List):
    final_txt=[]
    for idx in range(len(list)):
        res = list[idx]
        for line in res:
            final_txt.append(line[1][0])
    #print("\n".join(final_txt))
    return "\n".join(final_txt)